import cv2
import numpy as np
import mediapipe as mp
from PIL import Image, ImageDraw
import os
import logging
from typing import Dict, List, Tuple, Optional
import json

logger = logging.getLogger(__name__)

class AdvancedFaceDetector:
    """基于MediaPipe的高精度面部检测和关键点提取"""
    
    def __init__(self):
        self.mp_face_mesh = mp.solutions.face_mesh
        self.mp_drawing = mp.solutions.drawing_utils
        self.mp_drawing_styles = mp.solutions.drawing_styles
        
        # 初始化面部网格检测器
        self.face_mesh = self.mp_face_mesh.FaceMesh(
            static_image_mode=True,
            max_num_faces=1,
            refine_landmarks=True,
            min_detection_confidence=0.5
        )
        
        # 嘴部关键点索引 (MediaPipe 468点模型)
        self.MOUTH_LANDMARKS = {
            'outer_lip': [61, 84, 17, 314, 405, 320, 307, 375, 321, 308, 324, 318],
            'upper_lip': [13, 82, 81, 80, 78, 191, 80, 81, 82, 13, 312, 311, 310, 415],
            'lower_lip': [14, 87, 178, 88, 95, 146, 91, 181, 84, 17, 314, 405, 320, 307],
            'inner_lip': [78, 191, 80, 81, 82, 13, 312, 311, 310, 415, 308, 324, 318]
        }
        
        # 面部区域关键点
        self.FACE_OVAL = [10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288,
                         397, 365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136,
                         172, 58, 132, 93, 234, 127, 162, 21, 54, 103, 67, 109]

    def detect_face_landmarks(self, image_path: str) -> Optional[Dict]:
        """检测面部关键点"""
        try:
            # 读取图像
            image = cv2.imread(image_path)
            if image is None:
                logger.error(f"无法读取图像: {image_path}")
                return None
                
            # 转换为RGB
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            height, width = rgb_image.shape[:2]
            
            # 检测面部关键点
            results = self.face_mesh.process(rgb_image)
            
            if not results.multi_face_landmarks:
                logger.warning("未检测到面部")
                return None
                
            # 获取第一个面部的关键点
            face_landmarks = results.multi_face_landmarks[0]
            
            # 转换为像素坐标
            landmarks = []
            for landmark in face_landmarks.landmark:
                x = int(landmark.x * width)
                y = int(landmark.y * height)
                landmarks.append((x, y))
            
            return {
                'landmarks': landmarks,
                'image_size': (width, height),
                'mouth_regions': self._extract_mouth_regions(landmarks),
                'face_bbox': self._get_face_bbox(landmarks)
            }
            
        except Exception as e:
            logger.error(f"面部检测失败: {str(e)}")
            return None

    def _extract_mouth_regions(self, landmarks: List[Tuple[int, int]]) -> Dict:
        """提取嘴部区域坐标"""
        mouth_regions = {}
        
        for region_name, indices in self.MOUTH_LANDMARKS.items():
            region_points = [landmarks[i] for i in indices if i < len(landmarks)]
            if region_points:
                # 计算边界框
                xs = [p[0] for p in region_points]
                ys = [p[1] for p in region_points]
                mouth_regions[region_name] = {
                    'points': region_points,
                    'bbox': (min(xs), min(ys), max(xs), max(ys)),
                    'center': (sum(xs) // len(xs), sum(ys) // len(ys))
                }
        
        return mouth_regions

    def _get_face_bbox(self, landmarks: List[Tuple[int, int]]) -> Tuple[int, int, int, int]:
        """获取面部边界框"""
        face_points = [landmarks[i] for i in self.FACE_OVAL if i < len(landmarks)]
        if not face_points:
            # 回退到所有关键点
            face_points = landmarks
            
        xs = [p[0] for p in face_points]
        ys = [p[1] for p in face_points]
        return (min(xs), min(ys), max(xs), max(ys))

    def create_precise_mouth_variants(self, image_path: str, character_id: str, 
                                    output_dir: str) -> Dict[str, str]:
        """基于精确面部检测创建嘴型变体"""
        try:
            # 检测面部关键点
            face_data = self.detect_face_landmarks(image_path)
            if not face_data:
                logger.error("面部检测失败，使用默认方法")
                return self._fallback_mouth_variants(image_path, character_id, output_dir)
            
            # 加载原始图像
            original_image = Image.open(image_path).convert('RGBA')
            width, height = original_image.size
            
            # 获取嘴部区域
            mouth_region = face_data['mouth_regions'].get('outer_lip')
            if not mouth_region:
                logger.error("未能提取嘴部区域")
                return self._fallback_mouth_variants(image_path, character_id, output_dir)
            
            mouth_bbox = mouth_region['bbox']
            mouth_center = mouth_region['center']
            mouth_points = mouth_region['points']
            
            # 计算嘴部尺寸
            mouth_width = mouth_bbox[2] - mouth_bbox[0]
            mouth_height = mouth_bbox[3] - mouth_bbox[1]
            
            variants = {}
            
            # 创建不同的嘴型变体
            for variant_name in ['closed', 'half_open', 'open']:
                variant_image = original_image.copy()
                
                # 创建遮罩来覆盖原始嘴部
                mask = Image.new('RGBA', (width, height), (0, 0, 0, 0))
                mask_draw = ImageDraw.Draw(mask)
                
                # 用肤色遮盖原始嘴部（简化处理）
                mask_draw.polygon(mouth_points, fill=(255, 220, 177, 255))
                
                # 应用遮罩
                variant_image = Image.alpha_composite(variant_image, mask)
                
                # 绘制新的嘴型
                draw = ImageDraw.Draw(variant_image)
                
                if variant_name == 'closed':
                    # 闭嘴 - 细线
                    draw.line([
                        (mouth_bbox[0], mouth_center[1]),
                        (mouth_bbox[2], mouth_center[1])
                    ], fill=(139, 69, 19), width=2)
                    
                elif variant_name == 'half_open':
                    # 半开 - 小椭圆
                    oval_height = mouth_height // 3
                    draw.ellipse([
                        (mouth_bbox[0] + mouth_width // 6, mouth_center[1] - oval_height // 2),
                        (mouth_bbox[2] - mouth_width // 6, mouth_center[1] + oval_height // 2)
                    ], fill=(34, 34, 34), outline=(139, 69, 19), width=1)
                    
                elif variant_name == 'open':
                    # 大开 - 大椭圆
                    oval_height = mouth_height // 2
                    draw.ellipse([
                        (mouth_bbox[0], mouth_center[1] - oval_height // 2),
                        (mouth_bbox[2], mouth_center[1] + oval_height // 2)
                    ], fill=(34, 34, 34), outline=(139, 69, 19), width=1)
                
                # 保存变体
                variant_path = os.path.join(output_dir, f"{character_id}_{variant_name}.png")
                variant_image.convert('RGB').save(variant_path)
                variants[variant_name] = variant_path
                
                logger.info(f"创建嘴型变体: {variant_name} -> {variant_path}")
            
            # 保存面部数据用于调试
            debug_data = {
                'face_bbox': face_data['face_bbox'],
                'mouth_center': mouth_center,
                'mouth_bbox': mouth_bbox,
                'mouth_width': mouth_width,
                'mouth_height': mouth_height
            }
            
            debug_path = os.path.join(output_dir, f"{character_id}_face_data.json")
            with open(debug_path, 'w', encoding='utf-8') as f:
                json.dump(debug_data, f, ensure_ascii=False, indent=2)
            
            return variants
            
        except Exception as e:
            logger.error(f"精确嘴型变体创建失败: {str(e)}")
            return self._fallback_mouth_variants(image_path, character_id, output_dir)

    def _fallback_mouth_variants(self, image_path: str, character_id: str, 
                               output_dir: str) -> Dict[str, str]:
        """回退到简单的嘴型变体创建方法"""
        logger.info("使用回退方法创建嘴型变体")
        
        try:
            with Image.open(image_path) as img:
                width, height = img.size
                
                # 简化的面部估算
                face_center_x = width // 2
                face_center_y = int(height * 0.4)  # 调整面部中心位置
                
                # 嘴部位置
                mouth_x = face_center_x
                mouth_y = int(face_center_y * 1.3)
                
                # 嘴部尺寸
                mouth_width = width // 12
                mouth_height = height // 25
                
                variants = {}
                
                for variant in ['closed', 'half_open', 'open']:
                    img_variant = img.copy()
                    draw = ImageDraw.Draw(img_variant)
                    
                    mouth_left = mouth_x - mouth_width // 2
                    mouth_right = mouth_x + mouth_width // 2
                    
                    if variant == 'closed':
                        draw.line([(mouth_left, mouth_y), (mouth_right, mouth_y)], 
                                fill=(139, 69, 19), width=2)
                    elif variant == 'half_open':
                        draw.ellipse([
                            (mouth_left, mouth_y - mouth_height // 4),
                            (mouth_right, mouth_y + mouth_height // 4)
                        ], outline=(139, 69, 19), width=2)
                    elif variant == 'open':
                        draw.ellipse([
                            (mouth_left, mouth_y - mouth_height // 2),
                            (mouth_right, mouth_y + mouth_height // 2)
                        ], outline=(139, 69, 19), width=2)
                    
                    variant_path = os.path.join(output_dir, f"{character_id}_{variant}.png")
                    img_variant.save(variant_path)
                    variants[variant] = variant_path
                
                return variants
                
        except Exception as e:
            logger.error(f"回退方法也失败了: {str(e)}")
            return {}

    def visualize_face_detection(self, image_path: str, output_path: str) -> bool:
        """可视化面部检测结果，用于调试"""
        try:
            face_data = self.detect_face_landmarks(image_path)
            if not face_data:
                return False
            
            # 读取原始图像
            image = cv2.imread(image_path)
            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            # 绘制面部关键点
            for i, (x, y) in enumerate(face_data['landmarks']):
                cv2.circle(image_rgb, (x, y), 1, (0, 255, 0), -1)
            
            # 绘制嘴部区域
            for region_name, region_data in face_data['mouth_regions'].items():
                if region_name == 'outer_lip':
                    points = np.array(region_data['points'], dtype=np.int32)
                    cv2.polylines(image_rgb, [points], True, (255, 0, 0), 2)
                    
                    # 绘制边界框
                    bbox = region_data['bbox']
                    cv2.rectangle(image_rgb, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), 2)
                    
                    # 绘制中心点
                    center = region_data['center']
                    cv2.circle(image_rgb, center, 5, (255, 255, 0), -1)
            
            # 绘制面部边界框
            face_bbox = face_data['face_bbox']
            cv2.rectangle(image_rgb, (face_bbox[0], face_bbox[1]), 
                         (face_bbox[2], face_bbox[3]), (0, 255, 255), 2)
            
            # 保存结果
            result_image = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
            cv2.imwrite(output_path, result_image)
            
            logger.info(f"面部检测可视化保存至: {output_path}")
            return True
            
        except Exception as e:
            logger.error(f"可视化失败: {str(e)}")
            return False


# 集成到CharacterAnimator类
class EnhancedCharacterAnimator:
    """增强的角色动画器，使用高精度面部检测"""
    
    def __init__(self, api_key: str = None):
        self.api_key = api_key
        self.face_detector = AdvancedFaceDetector()
    
    def create_precise_mouth_variants(self, character_image_path: str, character_id: str) -> bool:
        """创建精确的嘴型变体"""
        try:
            output_dir = os.path.dirname(character_image_path)
            
            # 使用高精度面部检测创建变体
            variants = self.face_detector.create_precise_mouth_variants(
                character_image_path, character_id, output_dir
            )
            
            if not variants:
                logger.error("无法创建嘴型变体")
                return False
            
            # 创建调试可视化
            debug_path = os.path.join(output_dir, f"{character_id}_face_debug.png")
            self.face_detector.visualize_face_detection(character_image_path, debug_path)
            
            logger.info(f"成功创建 {len(variants)} 个嘴型变体")
            return True
            
        except Exception as e:
            logger.error(f"创建嘴型变体失败: {str(e)}")
            return False

    def has_mouth_variants(self, character_id: str, images_dir: str) -> bool:
        """检查是否存在嘴型变体"""
        variants = ["closed", "half_open", "open"]
        return all(
            os.path.exists(os.path.join(images_dir, f"{character_id}_{v}.png")) 
            for v in variants
        )


# 使用示例
if __name__ == "__main__":
    # 测试面部检测
    detector = AdvancedFaceDetector()
    
    # 检测面部关键点
    result = detector.detect_face_landmarks("character.png")
    if result:
        print("检测到面部关键点:", len(result['landmarks']))
        print("嘴部中心:", result['mouth_regions']['outer_lip']['center'])
    
    # 创建嘴型变体
    animator = EnhancedCharacterAnimator()
    animator.create_precise_mouth_variants("character.png", "test_character")
